#include "common.h"

// 1. RDMA连接相关的资源
static struct ibv_pd *pd = NULL;                // 保护域
static struct ibv_cq *cq = NULL;                // 完成队列

static struct ibv_qp_init_attr qp_attr;         // 创建QP的属性信息
static struct ibv_qp     *client_qp    = NULL;
static struct rdma_cm_id *cm_server_id = NULL, 
                         *cm_client_id = NULL;

// CM事件管道
// RDMA IO通信事件管道
static struct rdma_event_channel *cm_event_channel      = NULL;
static struct ibv_comp_channel   *io_comp_channel = NULL;

// RDMA操作的真实数据和元数据的MR信息
static struct ibv_mr *client_meta_mr    = NULL, 
					 *server_buffer_mr  = NULL, 
					 *server_meta_mr    = NULL;


static struct meta_attr client_meta_attr, server_meta_attr;

// 接收和发送的WR和SGE
static struct ibv_recv_wr client_recv_wr, *bad_client_recv_wr = NULL;
static struct ibv_send_wr server_send_wr, *bad_server_send_wr = NULL;
static struct ibv_sge client_recv_sge, server_send_sge;


/* Starts an RDMA server by allocating basic connection resources */
//! 1. 启动RDMA服务端并分配基本的资源
static int start(struct sockaddr_in *server_addr) 
{
	int ret = -1;
	struct rdma_cm_event *cm_event = NULL;

	// 1. 创建RDMA CM事件管道
	cm_event_channel = rdma_create_event_channel();
	if (!cm_event_channel) {
		rdma_error("创建服务端CM事件管道失败, 错误码 : (%d)", -errno);
		return -errno;
	}
	rdma_debug("创建服务端CM事件管道成功: %p \n", cm_event_channel);
	

	// 2. 创建服务端CM ID(类似于服务端套接字)
	ret = rdma_create_id(cm_event_channel, &cm_server_id, NULL, RDMA_PS_TCP);
	if (ret) {
		rdma_error("创建服务端CM ID失败, 错误码: %d ", -errno);
		return -errno;
	}
	rdma_debug("创建服务端CM ID成功 \n");



	// 3. 类似于绑定套接字
	ret = rdma_bind_addr(cm_server_id, (struct sockaddr*) server_addr);
	if (ret) {
		rdma_error("绑定服务端CM ID失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("绑定服务端CM ID成功 \n");


	// 4. 开始监听客户端接入
	ret = rdma_listen(cm_server_id, 8);  // 允许阻塞的请求队列长度(超过的请求可能会被拒绝)
	if (ret) {
		rdma_error("服务端监听失败, 错误码: %d ", -errno);
		return -errno;
	}
	printf("服务端成功监听在: %s , port: %d \n",
			inet_ntoa(server_addr->sin_addr),
			ntohs(server_addr->sin_port));
	

	// 5. 阻塞等待客户端连接
	ret = cm_event_callback(cm_event_channel, 
			RDMA_CM_EVENT_CONNECT_REQUEST,
			&cm_event);
	if (ret) {
		rdma_error("阻塞等待客户端连接失败, 错误码 %d \n" , ret);
		return ret;
	}


	// 有链接进来，获取客户端套接字信息
	cm_client_id = cm_event->id;          

	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("释放客户端连接CM失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("新的客户端连接进来, CM ID 信息已存入: %p\n", cm_client_id);
	return ret;
}




//! 2. 分配RDMA通信相关的资源
static int setup()
{
	int ret = -1;
	if(!cm_client_id){
		rdma_error("客户端CM ID信息为空, 退出... \n");
		return -EINVAL;
	}



	// 现在有了客户端套接字信息，开始分配RDMA通信相关的资源
	// 1. 创建保护域
	pd = ibv_alloc_pd(cm_client_id->verbs);
	if (!pd) {
		rdma_error("服务端创建保护域失败, 错误码: %d\n",
				-errno);
		return -errno;
	}
	rdma_debug("服务端成功创建保护域于: %p \n", pd);


	// 2. 创建IO CM事件通知管道
	io_comp_channel = ibv_create_comp_channel(cm_client_id->verbs);
	if (!io_comp_channel) {
		rdma_error("服务端创建IO CM事件通知管道失败, 错误码: %d\n",
				-errno);
		return -errno;
	}
	rdma_debug("服务端成功创建IO CM事件通知管道于: %p \n", io_comp_channel);

	

	// 3. 创建完成队列
	cq = ibv_create_cq(cm_client_id->verbs,    // RDMA设备上下文 
							CQ_CAPACITY,       // 完成队列的荣阿玲 
							NULL,              // 用户上下文
							io_comp_channel,   // IO完成事件管道
							0                  // 信号向量
						);
	if (!cq) {
		rdma_error("服务端CQ创建失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("服务端创建CQ成功: %p 持有 %d 个元素! \n", cq, cq->cqe);

	
	// int ibv_req_notify_cq(struct ibv_cq *cq, 
	//                      int solicited_only);   // 通知出发条件的标志
	//                      0: 通知所有事件
	//                      1: 只有被标记为"solicited"的工作完成才会触发通知
	ret = ibv_req_notify_cq(cq, 0);
	if (ret) {
		rdma_error("CQ通知条件设置失败, 错误码: %d\n", -errno);
		return -errno;
	}


	// 4. 创建QP
	// 4.1 初始化QP属性
	bzero(&qp_attr, sizeof qp_attr);	
	qp_attr.cap.max_recv_sge = MAX_SGE;    // 接收SGE容量
	qp_attr.cap.max_recv_wr  = MAX_WR;     // 接收WR容量
	qp_attr.cap.max_send_sge = MAX_SGE;    // 发送SGE容量
	qp_attr.cap.max_send_wr  = MAX_WR;     // 发送WR容量
	qp_attr.qp_type          = IBV_QPT_RC; // QP类型


	// 这里我们使用相同的完成队列, 也可以选择使用不同的队列
	qp_attr.recv_cq = cq; // Recv事件完成队列
	qp_attr.send_cq = cq; // Send事件完成队列

	// 4.2 创建QP
	ret = rdma_create_qp(cm_client_id,    // 指定QP的连接
		       pd,                      // 指定QP的保护域
		       &qp_attr);               // 指定QP的属性
	if (ret) {
		rdma_error("QP创建失败, 错误码: %d \n", -errno);
	    return -errno;
	}


	client_qp = cm_client_id->qp;
	rdma_debug("QP创建成功: %p \n", cm_client_id);

	return ret;
}




//! 3. 接收客户端连接
static int accept_conn()
{
	int ret = -1;
	struct rdma_conn_param conn_param;
	struct rdma_cm_event *cm_event = NULL;
	struct sockaddr_in client_sock; 


	if(!cm_client_id || !client_qp) {
		rdma_error("客户端资源未正确初始化...\n");
		return -EINVAL;
	}

	// 1. 创建客户端元数据MR
    client_meta_mr = mr_register(pd, &client_meta_attr,
								sizeof(client_meta_attr), 
								(IBV_ACCESS_LOCAL_WRITE));
	if(!client_meta_mr){
		rdma_error("注册客户端元数据MR失败!\n");
		return -ENOMEM;
	}



	// 2. 准备接收客户端元数据SGE
	client_recv_sge.addr   = (uint64_t)client_meta_mr->addr;
	client_recv_sge.length = (uint32_t)client_meta_mr->length;
	client_recv_sge.lkey   = (uint32_t)client_meta_mr->lkey;


	// 3. 准备接收客户端元数据WR
	bzero(&client_recv_wr, sizeof(client_recv_wr));
	client_recv_wr.sg_list = &client_recv_sge;
	client_recv_wr.num_sge = 1; // only one SGE

	// 4. 创建接收客户端元数的Post Recv WR
	ret = ibv_post_recv(client_qp, &client_recv_wr, &bad_client_recv_wr);
	if (ret) {
		rdma_error("客户端元数据接收WR创建失败, 错误码: %d \n", ret);
		return ret;
	}
	rdma_debug("客户端元数据接收WR创建成功! \n");



	// 5. 接收客户端连接
	memset(&conn_param, 0, sizeof(conn_param));
    conn_param.initiator_depth     = 3; 
    conn_param.responder_resources = 3;
    
	ret = rdma_accept(cm_client_id, &conn_param);
	if (ret) {
		rdma_error("接收客户端连接失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("等待客户端连接成功事件: RDMA_CM_EVENT_ESTABLISHED.\n");



	// 6. 连接成功，阻塞处理事件通知
    ret = cm_event_callback(cm_event_channel, RDMA_CM_EVENT_ESTABLISHED, &cm_event);
    if (ret) {
		rdma_error("获取客户端连接成功事件失败, 错误码: %d \n", -errno);
		return -errno;
	}


	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("释放客户端连接成功事件失败, 错误码: %d\n", -errno);
		return -errno;
	}


	memcpy(&client_sock,  rdma_get_peer_addr(cm_client_id), 
			sizeof(struct sockaddr_in));

	
	printf("[INFO]: 成功接收到来自客户端的请求: %s \n", 
			inet_ntoa(client_sock.sin_addr));


	return ret;
}





//! 4. 向客户端发送自己的元数据
static int send_meta() 
{
	int ret = -1;
	struct ibv_wc wc;


	// 1. 先等到接收客户端发送元数据
	ret = wc_event_callback(io_comp_channel, &wc, 1);
	if (ret != 1) {
		rdma_error("接收客户端元数据失败, 错误码: %d \n", ret);
		return ret;
	}
	printf("成功接收到客户端元数据信息，内容如下: \n");
	show_meta(&client_meta_attr);
	


	// 2. 准备发送服务端自己的元数据信息
	// 2.1 准备内容数据的MR
    server_buffer_mr = mr_alloc(pd, 
								client_meta_attr.length, 
								(IBV_ACCESS_LOCAL_WRITE|
								IBV_ACCESS_REMOTE_READ|
								IBV_ACCESS_REMOTE_WRITE));
	if(!server_buffer_mr){
		rdma_error("服务端内容数据MR分配失败!\n");
		return -ENOMEM;
	}



	// 2.2 准备元数据信息
	server_meta_attr.address  = (uint64_t)server_buffer_mr->addr;
	server_meta_attr.length   = (uint32_t)server_buffer_mr->length;
	server_meta_attr.key.lkey = (uint32_t)server_buffer_mr->lkey;

	// 2.3 准备服务端元数据的MR
	server_meta_mr = mr_register(pd, 
								&server_meta_attr, 
								sizeof(server_meta_attr),
								IBV_ACCESS_LOCAL_WRITE);
	if(!server_meta_mr){
		rdma_error("服务端元数据MR分配失败!\n");
		return -ENOMEM;
	}
      



	// 2.4 准备服务端元数据信息SGE
	server_send_sge.addr   = (uint64_t)&server_meta_attr;
	server_send_sge.length = (uint32_t)sizeof(server_meta_attr);
	server_send_sge.lkey   = (uint32_t)server_meta_mr->lkey;
	

	// 2.5 准备服务端元数据WR 
	bzero(&server_send_wr, sizeof(server_send_wr));
	server_send_wr.sg_list = &server_send_sge;
	server_send_wr.num_sge = 1; // only 1 SGE element in the array 
	server_send_wr.opcode = IBV_WR_SEND; // This is a send request 
	server_send_wr.send_flags = IBV_SEND_SIGNALED; // We want to get notification 


	// 2.6 提交服务端元数据Post Send WR
	ret = ibv_post_send(client_qp, &server_send_wr,  &bad_server_send_wr);
	if (ret) {
		rdma_error("服务端提交元数据发Post Send送请求失败, 错误码: %d \n",
				-errno);
		return -errno;
	}


	// 3. 检验发送请求结果
	ret = wc_event_callback(io_comp_channel, &wc, 1);
	if (ret != 1) {
		rdma_error("服务端发送元数据失败, 错误码: %d \n", ret);
		return ret;
	}
	rdma_debug("服务端发送元数据成功! \n");

	return 0;
}




//! 5. 清理RDMA链接资源
static int cleanup()
{
	int ret = -1;
	struct rdma_cm_event *cm_event = NULL;


	// 1. 等待客户端断开连接事件
	rdma_debug("服务端等待连接断事件: RDMA_CM_EVENT_DISCONNECTED\n");
	ret = cm_event_callback(cm_event_channel, 
			RDMA_CM_EVENT_DISCONNECTED, &cm_event);
	if (ret) {
		rdma_error("服务端等待连接断事件失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("释放连接断CM事件失败, 错误码: %d\n", -errno);
		return -errno;
	}
	printf("收到客户端断开连接事件... \n");




	// 2. 释放RDMA通信资源
	// 2.1 释放QP
	rdma_destroy_qp(cm_client_id);
	ret = rdma_destroy_id(cm_client_id);
	if (ret) {
		rdma_error("释放客户端ID失败, %d \n", -errno);
	}


	// 2.2 释放CQ
	ret = ibv_destroy_cq(cq);
	if (ret) {
		rdma_error("释放CQ失败, %d \n", -errno);
	}


	// 2.3 释放IO CM事件通知管道
	ret = ibv_destroy_comp_channel(io_comp_channel);
	if (ret) {
		rdma_error("释放IO CM事件通知管道失败, %d \n", -errno);
	}


	// 2.4 释放保MR
	mr_free(server_buffer_mr);
	mr_deregister(server_meta_mr);	
	mr_deregister(client_meta_mr);	


	// 2.5 释放保护域
	ret = ibv_dealloc_pd(pd);
	if (ret) {
		rdma_error("释放保护域失败, %d \n", -errno);
	}


	// 2.6 释放服务端CM ID
	ret = rdma_destroy_id(cm_server_id);
	if (ret) {
		rdma_error("释放服务端CM ID失败, %d \n", -errno);
	}

	// 3. 释放服务端CM事件管道
	rdma_destroy_event_channel(cm_event_channel);


	printf("服务端资源清理和服务关闭成功... \n");


	return 0;
}




void usage() 
{
	printf("Usage:\n");
	printf("rdma_server: [-a <server_addr>] [-p <server_port>]\n");
	printf("(default port is %d)\n", RDMA_PORT);
	exit(1);
}



int main(int argc, char **argv) 
{
	int ret, option;
	struct sockaddr_in server_sock;
	bzero(&server_sock, sizeof server_sock);

	server_sock.sin_family      = AF_INET; 		     // 默认IPV4
	server_sock.sin_addr.s_addr = htonl(INADDR_ANY); // 默认使用0.0.0.0


	// 解析命行参数
	while ((option = getopt(argc, argv, "a:p:")) != -1) {
		switch (option) {
			case 'a':   // 1. 服务端RDMA网卡IP 
				ret = get_addr(optarg, (struct sockaddr*) &server_sock);
				if (ret) {
					rdma_error("Invalid IP \n");
					 return ret;
				}
				break;
			case 'p':   // 2. 服务端RDMA端口
				server_sock.sin_port = htons(strtol(optarg, NULL, 0)); 
				break;
			default:
				usage();
				break;
		}
	}
	
	// 如果为传入服务端端口, 使用默认值
	if(!server_sock.sin_port) {
		server_sock.sin_port = htons(RDMA_PORT);
	}

	ret = start(&server_sock);
	if (ret) {
		rdma_error("启动服务端RDMA服务失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = setup();
	if (ret) { 
		rdma_error("创建RDMA连接资源失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = accept_conn();
	if (ret) {
		rdma_error("接收客户端连接失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = send_meta();
	if (ret) {
		rdma_error("向客户端发送服务端元数据失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = cleanup();
	if (ret) { 
		rdma_error("服务端清理RDMA资源失败: %d \n", ret);
		return ret;
	}

	return 0;
}
